package middleware

import (
	"encoding/json"
	"net/http"
	"time"

	"gitee.com/kessil/chicory/internal/models"
)

// ResponseWrapper 响应包装器
type ResponseWrapper struct {
	OriginUrl string      `json:"originUrl,omitempty"` // 来源URL
	Timestamp int64       `json:"timestamp,omitempty"` // 时间戳
	Data      interface{} `json:"data"`                // 实际数据
}

// JSONResponseWithOrigin 带来源信息的JSON响应中间件
func JSONResponseWithOrigin(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// 创建自定义的ResponseWriter来捕获响应
		customWriter := &responseCaptureWriter{
			ResponseWriter: w,
			request:        r,
		}
		next.ServeHTTP(customWriter, r)
	})
}

type responseCaptureWriter struct {
	http.ResponseWriter
	request      *http.Request
	responseBody []byte
	statusCode   int
}

func (w *responseCaptureWriter) Write(data []byte) (int, error) {
	w.responseBody = data
	return len(data), nil
}

func (w *responseCaptureWriter) WriteHeader(statusCode int) {
	w.statusCode = statusCode
}

func (w *responseCaptureWriter) flush() {
	// 解析原始响应
	var originalResponse models.Response
	if err := json.Unmarshal(w.responseBody, &originalResponse); err == nil {
		// 创建包装响应
		wrappedResponse := ResponseWrapper{
			OriginUrl: w.request.URL.String(),
			Timestamp: time.Now().Unix(),
			Data:      originalResponse,
		}

		// 序列化并发送
		wrappedData, _ := json.Marshal(wrappedResponse)
		w.ResponseWriter.Header().Set("Content-Type", "application/json")
		w.ResponseWriter.WriteHeader(w.statusCode)
		w.ResponseWriter.Write(wrappedData)
	} else {
		// 如果解析失败，直接发送原始响应
		w.ResponseWriter.Header().Set("Content-Type", "application/json")
		w.ResponseWriter.WriteHeader(w.statusCode)
		w.ResponseWriter.Write(w.responseBody)
	}
}
